import pandas as pd
import os

def analyze_splits():
    """分析训练集和测试集的分布"""
    
    print("=== 分析训练集和测试集分布 ===\n")
    
    # 读取临床+基因组数据的分割信息
    clinical_genomic_split = pd.read_csv('./raw/clinical+genomic_split.csv')
    
    print("1. clinical+genomic_split.csv 文件分析:")
    print(f"   总样本数: {len(clinical_genomic_split)}")
    print("   Split列分布:")
    split_counts = clinical_genomic_split['Split'].value_counts()
    for split, count in split_counts.items():
        print(f"     {split}: {count} 样本 ({count/len(clinical_genomic_split)*100:.1f}%)")
    
    print(f"\n   训练集样本: {split_counts.get('train', 0)}")
    print(f"   测试集样本: {split_counts.get('test', 0)}")
    
    # 显示一些训练集和测试集的样本
    print("\n   训练集前5个样本:")
    train_samples = clinical_genomic_split[clinical_genomic_split['Split'] == 'train'].head()
    for _, row in train_samples.iterrows():
        print(f"     {row['case_id']}")
    
    print("\n   测试集前5个样本:")
    test_samples = clinical_genomic_split[clinical_genomic_split['Split'] == 'test'].head()
    for _, row in test_samples.iterrows():
        print(f"     {row['case_id']}")
    
    # 检查其他可能包含分割信息的文件
    print("\n2. 检查其他文件中的分割信息:")
    
    # 检查CT患者文件
    if os.path.exists('./raw/CT_patientfiles.csv'):
        ct_patients = pd.read_csv('./raw/CT_patientfiles.csv')
        print(f"   CT_patientfiles.csv: {len(ct_patients)} 行")
        if 'Split' in ct_patients.columns:
            print("   包含Split列")
            ct_split_counts = ct_patients['Split'].value_counts()
            for split, count in ct_split_counts.items():
                print(f"     {split}: {count} 样本")
        else:
            print("   不包含Split列")
    
    # 检查MRI患者文件
    if os.path.exists('./raw/MRI_patientfiles.csv'):
        mri_patients = pd.read_csv('./raw/MRI_patientfiles.csv')
        print(f"   MRI_patientfiles.csv: {len(mri_patients)} 行")
        if 'Split' in mri_patients.columns:
            print("   包含Split列")
            mri_split_counts = mri_patients['Split'].value_counts()
            for split, count in mri_split_counts.items():
                print(f"     {split}: {count} 样本")
        else:
            print("   不包含Split列")
    
    # 检查WSI患者文件
    if os.path.exists('./raw/WSI_patientfiles.csv'):
        wsi_patients = pd.read_csv('./raw/WSI_patientfiles.csv')
        print(f"   WSI_patientfiles.csv: {len(wsi_patients)} 行")
        if 'Split' in wsi_patients.columns:
            print("   包含Split列")
            wsi_split_counts = wsi_patients['Split'].value_counts()
            for split, count in wsi_split_counts.items():
                print(f"     {split}: {count} 样本")
        else:
            print("   不包含Split列")
    
    # 检查带标签的CT和MRI文件
    print("\n3. 检查带标签的文件:")
    
    if os.path.exists('./raw/patients_with_labels_CT_final.csv'):
        ct_labeled = pd.read_csv('./raw/patients_with_labels_CT_final.csv')
        print(f"   patients_with_labels_CT_final.csv: {len(ct_labeled)} 行")
        if 'Split' in ct_labeled.columns:
            print("   包含Split列")
            ct_labeled_split_counts = ct_labeled['Split'].value_counts()
            for split, count in ct_labeled_split_counts.items():
                print(f"     {split}: {count} 样本")
        else:
            print("   不包含Split列")
    
    if os.path.exists('./raw/patients_with_labels_MR_final.csv'):
        mr_labeled = pd.read_csv('./raw/patients_with_labels_MR_final.csv')
        print(f"   patients_with_labels_MR_final.csv: {len(mr_labeled)} 行")
        if 'Split' in mr_labeled.columns:
            print("   包含Split列")
            mr_labeled_split_counts = mr_labeled['Split'].value_counts()
            for split, count in mr_labeled_split_counts.items():
                print(f"     {split}: {count} 样本")
        else:
            print("   不包含Split列")
    
    print("\n=== 总结 ===")
    print("1. clinical+genomic_split.csv 是主要的训练/测试集划分文件")
    print("2. 该文件包含了临床和基因组数据的完整分割信息")
    print("3. 其他文件可能包含特定模态（CT、MRI、WSI）的数据")
    print("4. 建议使用 clinical+genomic_split.csv 作为主要的分割参考")

if __name__ == '__main__':
    analyze_splits() 